Implement a MapSum class

Time: O(N); Space: O(T); med

Notes:

  • N is the length of key

  • T is the number of nodes in tree

Implement a MapSum class with insert, and sum methods.

For the method insert, you’ll be given a pair of (string, integer).

The string represents the key and the integer represents the value.

If the key already existed, then the original key-value pair will be overridden to the new one.

For the method sum, you’ll be given a string representing the prefix, and you need to return the sum of all the pairs’ value whose key starts with the prefix.

Example 1:

Input/Output:

  • insert(“apple”, 3) -> Null

  • sum(“ap”) -> 3

  • insert(“app”, 2) -> Null

  • sum(“ap”) -> 5

[1]:

import collections

class MapSum(object):

    def __init__(self):
        _trie = lambda: collections.defaultdict(_trie)
        self.__root = _trie()

    def insert(self, key, val):
        """
        :type key: str
        :type val: int
        :rtype: void
        """
        # Time: O(n)
        curr = self.__root
        for c in key:
            curr = curr[c]
        delta = val
        if "_end" in curr:
            delta -= curr["_end"]

        curr = self.__root
        for c in key:
            curr = curr[c]
            if "_count" in curr:
                curr["_count"] += delta
            else:
                curr["_count"] = delta
        curr["_end"] = val

    def sum(self, prefix):
        """
        :type prefix: str
        :rtype: int
        """
        # Time: O(n)
        curr = self.__root
        for c in prefix:
            if c not in curr:
                return 0
            curr = curr[c]
        return curr["_count"]

[2]:

# Your MapSum object will be instantiated and called as such
obj = MapSum()

obj.insert("apple", 3)
assert obj.sum("ap") == 3
obj.insert("app", 2)
assert obj.sum("ap") == 5